import os
import torch
import numpy as np
import pandas as pd
from torch.utils import data
from config import DefaultConfig


def data_split(data_folder, target_folder, train_scale=0.8, val_scale=0.1):
    print("\t* Begin split the dataSet...")
    total_names = os.listdir(data_folder)
    split_names = ['train', 'val', 'test']
    for split_name in split_names:
        # Create three folders(train,val,test) in the target directory
        split_path = os.path.join(target_folder,split_name)
        # Determine if the path exists
        if os.path.isdir(split_path):
            pass
        else:
            os.mkdir(split_path)
    for class_name in total_names:
        print("\t* Raw data :"+class_name)
        # Read dataSet
        current_class_data_path = os.path.join(data_folder, class_name)
        current_all_data = pd.read_csv(current_class_data_path, header=None)
        current_data_length = len(current_all_data)
        # Give file path and split-flag
        train_folder = os.path.join(os.path.join(target_folder, 'train'), 'train_'+class_name[:-4]+'.csv')
        val_folder = os.path.join(os.path.join(target_folder, 'val'), 'val_'+class_name[:-4]+'.csv')
        test_folder = os.path.join(os.path.join(target_folder, 'test'), 'test_'+class_name[:-4]+'.csv')
        train_stop_flag = int(current_data_length * train_scale)
        val_stop_flag = int(current_data_length * (train_scale+val_scale))
        # Save the data
        train = current_all_data.iloc[0:train_stop_flag,:]
        train.to_csv(train_folder,sep=',',header=0,index=0)
        val = current_all_data.iloc[train_stop_flag:val_stop_flag,:]
        val.to_csv(val_folder,sep=',',header=0,index=0)
        test = current_all_data.iloc[val_stop_flag:, :]
        test.to_csv(test_folder,sep=',',header=0,index=0)
    print("\t* Finish the split of dataSet...")


class ScatterData(data.Dataset):
    def __init__(self, x_path, y_path):
        super(ScatterData, self).__init__()
        self.data = pd.read_csv(x_path, header=None).values

        self.data_mean = self.data.mean(axis=0)
        self.data_std = self.data.std(axis=0)
        self.data = (self.data-self.data_mean)/self.data_std

        self.label = pd.read_csv(y_path, header=None).values
        self.data_num = self.data.shape[0]
        self.data_size = self.data.shape[1]
        self.label_size = self.label.shape[1]

    def __getitem__(self, index):
        return {'x': self.data[index, :], 'y': self.label[index, :]}

    def __len__(self):
        return self.data_num


if __name__ == '__main__':
    opt = DefaultConfig()
    data_split(opt.data_folder_path, opt.target_folder_path)
